import sqlite3
import csv
from collections import defaultdict

# 数据库连接
conn = sqlite3.connect("E:\\immunedb\\immuno.db")
cursor = conn.cursor()

try:
    # 指定需要分析的HLA类型
    target_hlas = ["HLA-A*02:01", "HLA-A*01:01", "HLA-B*48:01"]
    print(f"分析的HLA类型：{target_hlas}")

    # 步骤1：先获取所有符合条件的peptide_id对应的avg_exp，计算全局平均值
    print("\n=== 计算全局avg_exp平均值 ===")
    all_avg_exps = []
    for hla in target_hlas:  # 循环变量是小写hla
        # 获取该HLA中probability > 0.5的peptide_id对应的avg_exp
        cursor.execute("""
            SELECT DISTINCT pg.avg_exp
            FROM predictions p
            JOIN peptides pe ON p.peptide_id = pe.peptide_id
            JOIN protein_gene pg ON pe.protein_id = pg.protein_id
            WHERE p.hla = ? AND p.probability > 0.5
        """, (hla,))
        exp_values = [row[0] for row in cursor.fetchall() if row[0] is not None]
        all_avg_exps.extend(exp_values)
        print(f"  已收集{hla}的{len(exp_values)}个avg_exp值")  # 此处修正为小写hla

    if not all_avg_exps:
        raise ValueError("未找到任何avg_exp数据，无法计算平均值")
    
    # 计算全局avg_exp平均值
    global_avg_exp = sum(all_avg_exps) / len(all_avg_exps)
    print(f"  所有符合条件数据的avg_exp平均值为：{global_avg_exp:.4f}")

    # 步骤2：统计每个HLA的肽段在组织中的出现次数及名称（仅保留avg_exp > 全局平均值的数据）
    result = {
        hla: {
            'counts': defaultdict(int),
            'tissues': defaultdict(set)
        } for hla in target_hlas
    }

    for hla in target_hlas:
        print(f"\n处理HLA: {hla}")
        
        # 获取该HLA中probability > 0.5且avg_exp > 全局平均值的peptide_id
        cursor.execute("""
            SELECT DISTINCT p.peptide_id
            FROM predictions p
            JOIN peptides pe ON p.peptide_id = pe.peptide_id
            JOIN protein_gene pg ON pe.protein_id = pg.protein_id
            WHERE p.hla = ? 
              AND p.probability > 0.5 
              AND pg.avg_exp > ?
        """, (hla, global_avg_exp))
        peptide_ids = [row[0] for row in cursor.fetchall()]
        print(f"  找到{len(peptide_ids)}个符合条件（avg_exp > 全局平均值）的peptide_id")

        if not peptide_ids:
            continue

        # 遍历每个peptide_id，统计组织出现次数及名称
        for idx, peptide_id in enumerate(peptide_ids, 1):
            # 通过peptide_id关联protein_id
            cursor.execute("""
                SELECT protein_id 
                FROM peptides 
                WHERE peptide_id = ?
            """, (peptide_id,))
            protein_id = cursor.fetchone()
            if not protein_id:
                continue
            protein_id = protein_id[0]

            # 通过protein_id关联tissue和avg_exp（从protein_gene表）
            cursor.execute("""
                SELECT DISTINCT pg.tissue, pg.avg_exp
                FROM protein_gene pg
                WHERE pg.protein_id = ?
                  AND pg.avg_exp > ?
            """, (protein_id, global_avg_exp))
            tissue_data = cursor.fetchall()
            if not tissue_data:
                continue

            # 提取组织名称（去重后）
            tissues = list({row[0] for row in tissue_data})
            tissue_count = len(tissues)

            # 更新统计结果
            result[hla]['counts'][tissue_count] += 1
            for tissue in tissues:
                result[hla]['tissues'][tissue_count].add(tissue)
            
            # 进度提示
            if idx % 100 == 0:
                print(f"  已处理{idx}/{len(peptide_ids)}个肽段")

    # 步骤3：整理结果并写入CSV
    output_path = "E:/immunedb/hla_high_avg_exp_tissue_details.csv"
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "HLA类型", 
            "组织出现次数", 
            "肽段数量", 
            "涉及的组织名称",
            "全局avg_exp平均值（筛选阈值）"
        ])
        
        for hla in target_hlas:
            tissue_counts = sorted(result[hla]['counts'].keys())
            for count in tissue_counts:
                tissue_names = ",".join(sorted(result[hla]['tissues'][count]))
                writer.writerow([
                    hla,
                    count,
                    result[hla]['counts'][count],
                    tissue_names,
                    round(global_avg_exp, 4)
                ])
            # 补充1-3次的记录
            for count in range(1, 4):
                if count not in result[hla]['counts']:
                    writer.writerow([
                        hla, 
                        count, 
                        0, 
                        "",
                        round(global_avg_exp, 4)
                    ])

    print(f"\n统计结果已保存到：{output_path}")
    print("\n统计摘要：")
    for hla in target_hlas:
        print(f"\n{hla}：")  # 此处也修正为小写hla
        for count in sorted(result[hla]['counts'].keys()):
            tissues = sorted(result[hla]['tissues'][count])
            print(f"  出现{count}次的肽段数量：{result[hla]['counts'][count]}，涉及组织：{tissues}")
    print(f"\n全局avg_exp筛选阈值：{global_avg_exp:.4f}")

except sqlite3.Error as e:
    print(f"\n数据库错误: {e}")
except Exception as e:
    print(f"\n其他错误: {e}")
finally:
    cursor.close()
    conn.close()
    print("\n程序结束，数据库连接已关闭")